#include "NeutrinoTools.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TLegend.h"
#include "TStyle.h"
#include "TDirectory.h"
#include "TSystem.h"

#include "style.h"

const Int_t gkdpT = 5;
const Int_t gkdphiT = 6;
const Int_t gkdalphaT = 7;
const Int_t gkProtonPTT = 8;
const Int_t gkMuonKT = 9;

Bool_t ToReturn(const Int_t kSigmaMode, 
                const Int_t nhitcut, const Double_t larm, const Double_t maxp, const Double_t maxchi,
                const Double_t fpath, const Int_t maxnhit, 
                const Bool_t kZaxis, const Bool_t kTrueMu, const Int_t kNuDirOpt)
{
  const TString pwd=gSystem->pwd();
  
  if(pwd.Contains("SK")){
    if(kZaxis || kTrueMu || kNuDirOpt!=1 ){
      return kTRUE;
    }
    
    if(nhitcut || larm>0 || maxp>0 || maxchi>0  || fpath>0 || maxnhit>0 ){
      return kTRUE;
    }      
  }
  else{
    if(kZaxis){
      //require PID, others not care
      if(kNuDirOpt!=2){//kTrueMu==0){
        return kTRUE;
      }
    }
    
   //f(kTrueMu){
   // if(kNuDirOpt!=2){
   //   return kTRUE;
   // }
   //
  }
  
  if((kNuDirOpt== gkdpT||kNuDirOpt== gkProtonPTT||kNuDirOpt== gkMuonKT) && kSigmaMode!=0){
    return kTRUE;
  }

  if(kNuDirOpt== gkdphiT && kSigmaMode!=2){
    return kTRUE;
  }

  if(kNuDirOpt== gkdalphaT && kSigmaMode!=2){
    return kTRUE;
  }

  return kFALSE;
}

TCanvas * getCanvas()
{
  gStyle->SetOptStat(0);
  
  style::SetGlobalStyle();

  style::fgkYTitleOffset = 1.4;
  style::fgkXTitleOffset = 1.4;

  TCanvas *c0= new TCanvas;
  style::PadSetup(c0);

  const Double_t lm=0.15;
  const Double_t rm=0.13;
  const Double_t tm=0.13;
  const Double_t bm=0.17;

  c0->SetLeftMargin(lm);
  c0->SetRightMargin(rm);
  c0->SetTopMargin(tm);
  c0->SetBottomMargin(bm);

  gStyle->SetOptTitle(1);
  //gStyle->SetTitleX(0.05);
  gStyle->SetTitleW(0.99);

  return c0;
}

void setHist(const Int_t kNuDirOpt, TH2D *& h2, TH1D *& halleff, TH1D *& heff, TH1D *& hallmuonpurity, TH1D *& hmuonpurity, TH1D *& hallnupurity, TH1D *& hnupurity)
{
  const Int_t xnbin = 20;
  Double_t xmin = 0.05;
  if( (kNuDirOpt==gkdpT||kNuDirOpt== gkProtonPTT||kNuDirOpt== gkMuonKT) || kNuDirOpt==gkdalphaT){
    xmin = 0;
  }
  if(kNuDirOpt==gkdphiT){
    xmin = 1e-2;
  }
  Double_t xmax = 2;
  if(kNuDirOpt==gkdpT||kNuDirOpt== gkProtonPTT||kNuDirOpt== gkMuonKT){
    xmax = 0.75;
  }
  if(kNuDirOpt==gkdphiT || kNuDirOpt==gkdalphaT){
    xmax = 3.145;
  }

  const Int_t ynbin = 120;//(kSigmaMode?1:2);
  //const Double_t ymax = (kSigmaMode? 0.6 : 0.3)*(pwd.Contains("SK")?1:2);
  const Double_t ymax = 1.2;//(kSigmaMode? 0.6 : 0.3)*2;
  const Double_t ymin = -ymax;
  //TH2D * h2=new TH2D("h2",cut, pwd.Contains("SK")?50:10,  pwd.Contains("SK")?0.1:0.2, 2, ynbin, ymin, ymax); NeutrinoTools::BinLog(h2->GetXaxis());
  h2=new TH2D("h2","", xnbin, xmin, xmax, ynbin, ymin, ymax); 

  halleff = new TH1D("halleff","",xnbin, xmin, xmax); 
  heff = new TH1D("heff","",xnbin, xmin, xmax); 

  hallmuonpurity = new TH1D("hallmuonpurity","",xnbin, xmin, xmax); 
  hmuonpurity = new TH1D("hmuonpurity","",xnbin, xmin, xmax); 

  hallnupurity = new TH1D("hallnupurity","",xnbin, xmin, xmax); 
  hnupurity = new TH1D("hnupurity","",xnbin, xmin, xmax); 

  if(xmin!=0){
    NeutrinoTools::BinLog(h2->GetXaxis());
    NeutrinoTools::BinLog(halleff->GetXaxis());
    NeutrinoTools::BinLog(heff->GetXaxis());
    NeutrinoTools::BinLog(hallmuonpurity->GetXaxis());
    NeutrinoTools::BinLog(hmuonpurity->GetXaxis());
    NeutrinoTools::BinLog(hallnupurity->GetXaxis());
    NeutrinoTools::BinLog(hnupurity->GetXaxis());
  }
}

void Config(TString & tag, TString & cut, TString & xRec, TString & xSim, TString & plvar, TString & xTitRec, TString & xTitSim, TString & tottitle, 
            const TString formula, const Int_t kSigmaMode, 
            const Int_t nhitcut, const Double_t larm, const Double_t maxp, const Double_t maxchi,
            const Double_t fpath, const Int_t maxnhit, 
            const Bool_t kZaxis, const Bool_t kTrueMu, const Int_t kNuDirOpt, const Bool_t kproton)
{
  const TString tag0=Form("kZaxis%d_kTrueMu%d_kNuDirOpt%d", kZaxis, kTrueMu, kNuDirOpt);
  tag=Form("nhit%02d_larm%03.0f_maxp%03.0f_maxchi%03.0f_fpath%03.0f_maxnhit%d_formula%s_kSigmaMode%d_kproton%d_%s",nhitcut, larm*100, maxp*100, maxchi*100, fpath*100, maxnhit, formula.Data(),  kSigmaMode, kproton, tag0.Data());

  //==================================
  const TString tmpcut0 =  "f@NRec>0";
  cut= tmpcut0
    + (nhitcut>0? Form(" && f@RecNhits[0]>%d", nhitcut):"") 
    + (larm>0?    Form(" && f@RecFlightPath->Pt()>%.2f", larm):"")
    + (maxp>0?    Form(" && f@Rec->P()<%.2f",maxp):"")
    + (maxchi>0?  Form(" && %.2f<(f@RecChi2/f@RecNDOF) && (f@RecChi2/f@RecNDOF)<%.2f", 1/maxchi, maxchi):"")
    + (fpath>0?   Form(" && f@RecFlightPath->Mag()>%.2f", fpath):"")
    + (maxnhit>0? Form(" && f@RecNhits[0]<%d", maxnhit):"")
    + (kTrueMu ?  " && fProtonSimPDG==2212 && fMuonSimPDG==13" : "");

  //===========================

  if(kNuDirOpt==0){  
    //pt w.r.t. (001) 
    xRec = "f@Rec->Pt()";
  }
  else if(kNuDirOpt==1){
    //using parent dec point - vertex
    xRec = "f@RecPt->Mag()";
  }
  else if(kNuDirOpt==2){
    //using true nutrino direction
    xRec = "f@PtGuess->Mag()";
  }
  else if(!(kNuDirOpt== gkdpT||kNuDirOpt== gkProtonPTT||kNuDirOpt== gkMuonKT) && kNuDirOpt!=gkdphiT && kNuDirOpt!=gkdalphaT){
    printf("\n!!wrong kNuDirOpt %d\n\n", kNuDirOpt);exit(1);
  }
  //true pt
  xSim = "f@SimPt->Mag()";

  //overwrite kNuDirOpt
  if(kZaxis){
    xRec = "f@Rec->Pt()";
    xSim = "f@Sim->Pt()";
  }

  if(kNuDirOpt== gkdpT){
    xRec = "fDeltaRecPt->Mag()";
    xSim = "fDeltaSimPt->Mag()";
  }

  if(kNuDirOpt== gkdphiT){
    xRec = "fDeltaRecPt->Phi()";
    xSim = "fDeltaSimPt->Phi()";
  }

  if(kNuDirOpt== gkdalphaT){
    xRec = "fDeltaRecPt->Theta()";
    xSim = "fDeltaSimPt->Theta()";
  }

  if(kNuDirOpt== gkProtonPTT){
    xRec = "fProtonPTTRec";
    xSim = "fProtonPTTSim";
  }

  if(kNuDirOpt== gkMuonKT){
    xRec = "fMuonKTRec";
    xSim = "fMuonKTSim";
  }

  if(kSigmaMode==2){
    plvar=Form("%s:%s>>h2","rec-sim","sim");
  }
  else if(kSigmaMode==1){
    plvar=Form("%s:%s>>h2","1/rec-1/sim","sim");
  }
  else if(kSigmaMode==0){
    plvar=Form("%s:%s>>h2","rec/sim-1","sim");
  }
  
  plvar.ReplaceAll("rec",xRec);
  plvar.ReplaceAll("sim",xSim);

  //===========================

  if(kSigmaMode==2){
    tottitle = cut+Form(" %s;%s;%s", tag0.Data(), "$", "rrr-sss");
  }
  else if(kSigmaMode==1){
    tottitle = cut+Form(" %s;%s;%s", tag0.Data(), "$", "1/rrr-1/sss (c/GeV)");
  }
  else if(kSigmaMode==0){
    tottitle = cut+Form(" %s;%s;%s", tag0.Data(), "$", "rrr/sss-1");
  }

  tottitle.ReplaceAll("&&"," ");
  tottitle.ReplaceAll("   "," ");
  tottitle.ReplaceAll("  "," ");
  tottitle.ReplaceAll("f@NSim>0","");
  tottitle.ReplaceAll("f@NRec>0","");
  tottitle.ReplaceAll("f@RecNhits[0]","tpcnhit0");
  tottitle.ReplaceAll("f@RecNhits[1]","tpcnhit1");
  tottitle.ReplaceAll("f@RecFlightPath->Pt()","larm");
  tottitle.ReplaceAll("f@Charge[0]","tpccharge0");
  tottitle.ReplaceAll("f@Charge[1]","tpccharge1");
  tottitle.ReplaceAll("fMuonSimPDG==13","#mu^{-}");
  tottitle.ReplaceAll("fProtonSimPDG==2212","p");
  tottitle.ReplaceAll("f@Rec->P()","p_{#mu}");
  tottitle.ReplaceAll("f@RecFlightPath->Mag()","fpath");
  tottitle.ReplaceAll("(f@RecChi2/f@RecNDOF) (f@RecChi2/f@RecNDOF)","#chi^{2}_{track}/NDOF");

  //tottitle.ReplaceAll("sss", "p_{t}^{@ sim}");
  //tottitle.ReplaceAll("rrr", "p_{t}^{@ rec}");

  tottitle.ReplaceAll("sss", "|");
  tottitle.ReplaceAll("rrr", "~");

  tottitle.ReplaceAll("kZaxis0","");
  tottitle.ReplaceAll("kZaxis1_kTrueMu1_kNuDirOpt2","intrinsic p_{t}");
  
  tottitle.ReplaceAll("_k","k");

  tottitle.ReplaceAll("kTrueMu0","");
  tottitle.ReplaceAll("kTrueMu1","");
  
  tottitle.ReplaceAll("kNuDirOpt0","#vec{#nu}=(001)");
  tottitle.ReplaceAll("kNuDirOpt1","#vec{#nu}=D-V");
  tottitle.ReplaceAll(Form("kNuDirOpt%d",gkdpT),"#vec{#nu}=D-V");
  tottitle.ReplaceAll(Form("kNuDirOpt%d",gkdphiT),"#vec{#nu}=D-V");
  tottitle.ReplaceAll(Form("kNuDirOpt%d",gkdalphaT),"#vec{#nu}=D-V");
  tottitle.ReplaceAll(Form("kNuDirOpt%d",gkProtonPTT),"#vec{#nu}=D-V");
  tottitle.ReplaceAll(Form("kNuDirOpt%d",gkMuonKT),"#vec{#nu}=D-V");
  tottitle.ReplaceAll("kNuDirOpt2","#vec{#nu}_{sim}");

  tottitle.ReplaceAll("kZaxis1#vec{#nu}_{sim}","intrinsic p_{t}");

  //===

  if(kproton){
    cut.ReplaceAll("@","Proton");
    xRec.ReplaceAll("@","Proton");
    xSim.ReplaceAll("@","Proton");
    plvar.ReplaceAll("@","Proton");
    tottitle.ReplaceAll("@","p");
  }
  else{
    cut.ReplaceAll("@","Muon");
    xRec.ReplaceAll("@","Muon");
    xSim.ReplaceAll("@","Muon");
    plvar.ReplaceAll("@","Muon");
    tottitle.ReplaceAll("@","#mu");
  }

  //===

  xTitRec = NeutrinoTools::GetTitleFromVar(xRec);
  xTitSim = NeutrinoTools::GetTitleFromVar(xSim);
  tottitle.ReplaceAll("$", xTitSim);
  tottitle.ReplaceAll("|", xTitSim(0,xTitSim.First('(')));
  tottitle.ReplaceAll("~", xTitRec(0,xTitRec.First('(')));
}

void plotPtInvRes(const TString formula, const Int_t kSigmaMode=0, 
                  const Int_t nhitcut=0, const Double_t larm = 0, const Double_t maxp = 0, const Double_t maxchi = 0,
                  const Double_t fpath = 0, const Int_t maxnhit = 0, 
                  const Bool_t kZaxis=kFALSE, const Bool_t kTrueMu = kFALSE, const Int_t kNuDirOpt = 0, const Bool_t kproton = 0
                  )
{
  if(ToReturn(kSigmaMode, nhitcut, larm, maxp, maxchi, fpath, maxnhit, kZaxis, kTrueMu, kNuDirOpt)){
    return;
  }
  
  TTree * tree = (TTree*) gDirectory->Get("tree");

  //================================================
  TCanvas *c0 = getCanvas();

  TH2D * h2=0x0;
  TH1D * halleff=0x0;
  TH1D * heff = 0x0;
  TH1D * hallmuonpurity = 0x0;
  TH1D * hmuonpurity = 0x0;
  TH1D * hallnupurity = 0x0;
  TH1D * hnupurity = 0x0;

  setHist(kNuDirOpt, h2, halleff, heff, hallmuonpurity, hmuonpurity, hallnupurity, hnupurity);

  //================================================
  TString tag, cut, xRec, xSim, plvar, xTitRec, xTitSim, tottitle;
  Config(tag, cut, xRec, xSim, plvar, xTitRec, xTitSim, tottitle, formula, kSigmaMode, nhitcut, larm, maxp, maxchi, fpath, maxnhit, kZaxis, kTrueMu, kNuDirOpt, kproton);
  printf("check tag %s  || cut { %s } || xRec %s || xSim %s || plvar %s || xTitRec %s || xTitSim %s || tottitle (%s)\n\n",tag. Data(), cut.Data(), xRec.Data(), xSim.Data(), plvar.Data(), xTitRec.Data(), xTitSim.Data(), tottitle.Data() );

  tree->Draw(plvar,cut+ Form("&& f%sNSim>0", kproton?"Proton":"Muon"));

  tree->Draw(xSim+">>halleff");
  tree->Draw(xSim+">>heff",cut);

  tree->Draw(xRec+">>hallmuonpurity", cut);
  tree->Draw(xRec+">>hmuonpurity", cut+ (kproton? " && fProtonSimPDG==2212": " && fMuonSimPDG==13"));

  tree->Draw(xRec+">>hallnupurity", cut);
  tree->Draw(xRec+">>hnupurity", cut+" && fNeutrinoType==14 && fNeutMode<30 && fNeutMode>0");

  //============================================================================
  const Bool_t klogx = (kNuDirOpt!= gkdpT && kNuDirOpt!= gkProtonPTT && kNuDirOpt!= gkMuonKT && kNuDirOpt!= gkdalphaT); 

  h2->SetTitle(tottitle);

  //printf("check overflow and underflow: %f %f\n", h2->Integral(0,h2->GetNbinsX(), 0, 0), h2->Integral(0,h2->GetNbinsX(), h2->GetNbinsY(), h2->GetNbinsY()));

  TH2D * h2nor=NeutrinoTools::NormalHist(h2, 5, 1);

  gPad->SetLogx(klogx);
  style::ResetStyle(h2nor);
  h2nor->Draw("colz");
  
  c0->Print(Form("outplot/ptinvres2d%s.png", tag.Data()));

  gPad->SetLogz();
  c0->Print(Form("outplot/ptinvres2d%slogz.png", tag.Data()));

  //===========================
  const Double_t thres = 50;
  TList * lout=new TList;
  lout->Add(h2);

  TH1D *hrmsnor=0x0, *hrmsmpv=0x0, *hrmswid=0x0, *hrmsres=0x0, *hrmschi=0x0;
  NeutrinoTools::FitSlicesY(h2, hrmsnor, hrmsmpv, hrmswid, hrmsres, hrmschi, "RMS", thres, lout);

  TH1D *hnor=0x0, *hmpv=0x0, *hwid=0x0, *hres=0x0, *hchi=0x0;
  NeutrinoTools::FitSlicesY(h2, hnor, hmpv, hwid, hres, hchi, formula, thres, lout);

//  h2->FitSlicesY();
//  hmpv=h2_1;
//  hwid=h2_2;
  
  TFile * fout=new TFile(Form("outplot/fitslice%s.root", tag.Data()),"recreate");
  lout->Write();
  fout->Save();
  fout->Close();
  delete fout;

  //============================================================================
  //============================================================================
  
  const Double_t lm=0.15;
  const Double_t rm=0.02;
  const Double_t tm=0.13;
  const Double_t bm=0.17;

  TCanvas *c1=new TCanvas;
  style::PadSetup(c1);

  gPad->SetLogx(klogx);

  c1->SetLeftMargin(lm);
  c1->SetRightMargin(rm);
  c1->SetTopMargin(tm);
  c1->SetBottomMargin(bm);

  gPad->SetGrid();
  hmpv->SetTitle(tottitle);
  hmpv->SetYTitle(Form("<%s>", h2->GetYaxis()->GetTitle()));
  hmpv->SetMinimum(-0.2);
  hmpv->SetMaximum(0.2);
  style::ResetStyle(hmpv);

  hmpv->SetMarkerStyle(20);
  hmpv->SetMarkerSize(1);
  hmpv->SetMarkerColor(kBlack);
  hmpv->SetLineWidth(1);
  hmpv->SetLineColor(kBlack);
  hmpv->Draw("e");
  c1->Print(Form("outplot/ptinvresmean%s.png", tag.Data()));

  //============================================================================

  TCanvas *c2 =new TCanvas;
  style::PadSetup(c2);

  c2->SetLeftMargin(lm);
  c2->SetRightMargin(rm);
  c2->SetTopMargin(tm);
  c2->SetBottomMargin(bm);

  gPad->SetLogx(klogx);
  gPad->SetGrid();

  hwid->SetTitle(tottitle);
  hwid->SetYTitle(Form("#sigma(%s)", h2->GetYaxis()->GetTitle()));
  hwid->SetMinimum(0);
  //hwid->SetMaximum(pwd.Contains("SK")?0.2:0.4);//0.6
  if(kSigmaMode==0){
    hwid->SetMaximum(kNuDirOpt==1?0.5:1);
  }
  else if(kSigmaMode==1){
    hwid->SetMaximum(0.4);
  }
  else if(kSigmaMode==2){
    hwid->SetMaximum(0.7);
  }

  if(kNuDirOpt==gkdphiT){
    hwid->SetMaximum(kSigmaMode==2?0.1:0.7);
  }
  else if(kNuDirOpt==gkMuonKT){
    hwid->SetMaximum(0.3);
  }
  else if(kNuDirOpt==gkProtonPTT){
    hwid->SetMaximum(0.3);
  }
  else if(kNuDirOpt==gkdpT){
    hwid->SetMaximum(0.5);
  }
  else if(kNuDirOpt==gkdalphaT){
    hwid->SetMaximum(0.5);
  }

  style::ResetStyle(hwid);
  style::ResetStyle(hrmswid);
  
  hwid->SetMarkerStyle(20);
  hwid->SetMarkerSize(1);
  hwid->SetMarkerColor(kBlack);
  hwid->SetLineWidth(1);
  hwid->SetLineColor(kBlack);

  hrmswid->SetMarkerStyle(22);
  hrmswid->SetMarkerSize(1.5);
  hrmswid->SetMarkerColor(kRed);
  hrmswid->SetLineColor(kRed);
  hrmswid->SetLineWidth(1);

  hwid->Draw("e");
  hrmswid->Draw("same e");

  TLegend * wl = new TLegend(0.7,0.7,0.5+0.95,0.87);
  style::ResetStyle(wl,0.2);
  wl->SetTextAlign(12);
  //wl->SetBorderSize(1);
  wl->AddEntry(hrmswid, "RMS","lp");
  wl->AddEntry(hwid,formula,"lp");
  wl->Draw();

  c2->Print(Form("outplot/ptinvreswid%s.png", tag.Data()));

  if(  (kNuDirOpt!=1 && kNuDirOpt<gkdpT) || kZaxis!=0){
    return;
  }
  //============================================================================

  TCanvas *c3 =new TCanvas;
  style::PadSetup(c3);

  c3->SetLeftMargin(lm);
  c3->SetRightMargin(rm);
  c3->SetTopMargin(tm);
  c3->SetBottomMargin(bm);

  gPad->SetLogx(klogx);
  gPad->SetGrid();

  halleff->Sumw2();
  heff->Sumw2();
  heff->Divide(heff, halleff,1,1,"B");

  heff->SetTitle(Form(" %s;%s; efficiency",hwid->GetTitle(), xTitSim.Data()));
  heff->SetMinimum(gPad->GetLogy()?0.01:0);
  heff->SetMaximum(heff->GetBinContent(heff->GetMaximumBin())<0.4?0.5:1.1);
  style::ResetStyle(heff);

  heff->SetMarkerStyle(20);
  heff->SetMarkerSize(1);
  heff->SetMarkerColor(kBlack);
  heff->SetLineWidth(1);
  heff->SetLineColor(kBlack);
  heff->Draw("e");
  c3->Print(Form("outplot/ptinvreseff%s.png", tag.Data()));

  //============================================================================

  TCanvas *c4 =new TCanvas;
  style::PadSetup(c4);

  c4->SetLeftMargin(lm);
  c4->SetRightMargin(rm);
  c4->SetTopMargin(tm);
  c4->SetBottomMargin(bm);

  gPad->SetLogx(klogx);
  gPad->SetGrid();

  hallmuonpurity->Sumw2();
  hmuonpurity->Sumw2();
  hmuonpurity->Divide(hmuonpurity, hallmuonpurity,1,1,"B");

  //it is rec!!
  hmuonpurity->SetTitle(Form(" %s;%s; #mu^{-} purity",hwid->GetTitle(), xTitRec.Data()));
  hmuonpurity->SetMinimum(0.7);
  hmuonpurity->SetMaximum(1.1);
  style::ResetStyle(hmuonpurity);

  hmuonpurity->SetMarkerStyle(20);
  hmuonpurity->SetMarkerSize(1);
  hmuonpurity->SetMarkerColor(kBlack);
  hmuonpurity->SetLineWidth(1);
  hmuonpurity->SetLineColor(kBlack);
  hmuonpurity->Draw("e");
  c4->Print(Form("outplot/ptinvresmuonpurity%s.png", tag.Data()));

  //============================================================================

  TCanvas *c5 =new TCanvas;
  style::PadSetup(c5);

  c5->SetLeftMargin(lm);
  c5->SetRightMargin(rm);
  c5->SetTopMargin(tm);
  c5->SetBottomMargin(bm);

  gPad->SetLogx(klogx);
  gPad->SetGrid();

  hallnupurity->Sumw2();
  hnupurity->Sumw2();
  hnupurity->Divide(hnupurity, hallnupurity,1,1,"B");

  //it is rec!!
  hnupurity->SetTitle(Form(" %s;%s; #nu_{#mu} CC purity",hwid->GetTitle(), xTitRec.Data()));
  hnupurity->SetMinimum(0.7);
  hnupurity->SetMaximum(1.1);
  style::ResetStyle(hnupurity);

  hnupurity->SetMarkerStyle(20);
  hnupurity->SetMarkerSize(1);
  hnupurity->SetMarkerColor(kBlack);
  hnupurity->SetLineWidth(1);
  hnupurity->SetLineColor(kBlack);
  hnupurity->Draw("e");
  c5->Print(Form("outplot/ptinvresnupurity%s.png", tag.Data()));
}
